package geppetto.decoder;

import geppetto.decoder.scores.NumberOfOptionsScore;

import java.util.HashMap;
import java.util.LinkedList;

public class TranslationDecoder {
	protected HashMap<Integer, LinkedList<TranslationOption>> optionsByTargetPosition;
	protected int sourceSize;
	protected int targetSize;
	protected HashMap<String,TranslationHypothesisScore> scores;

	public TranslationDecoder(LinkedList<TranslationOption> options, int sourceSize, int targetSize) {
		super();
		this.sourceSize = sourceSize;
		this.targetSize = targetSize;
		this.optionsByTargetPosition = new HashMap<Integer, LinkedList<TranslationOption>>();
		for(TranslationOption option : options){
			getOptionsForTargetPosition(option.getTargetS()).add(option);
		}
		scores = new HashMap<String, TranslationHypothesisScore>();
	}
	
	protected LinkedList<TranslationOption> getOptionsForTargetPosition(int pos){
		if(!optionsByTargetPosition.containsKey(pos)){
			optionsByTargetPosition.put(pos, new LinkedList<TranslationOption>());
		}
		return optionsByTargetPosition.get(pos);
	}
	
	public LinkedList<TranslationHypothesis> getTranslationHypothesis() {
		LinkedList<TranslationHypothesis> finalHypothesis = new LinkedList<TranslationHypothesis>();
		LinkedList<TranslationHypothesis> stack = new LinkedList<TranslationHypothesis>();
		TranslationHypothesis nullHyp = TranslationHypothesis.generateNullHypothesis(sourceSize);
		for(TranslationHypothesisScore scorer : scores.values()){
			scorer.init(nullHyp);
		}
		stack.addFirst(nullHyp);
		while(!stack.isEmpty()){
			TranslationHypothesis root = stack.removeFirst();
			if(isHypothesisComplete(root)){
				finalHypothesis.add(root);
			}
			else{
				LinkedList<TranslationOption> possibleTargetOptions = getOptionsForTargetPosition(root.getTranslatedWordIndex()+1);
				LinkedList<TranslationOption> possibleOptions = new LinkedList<TranslationOption>();
				for(TranslationOption option : possibleTargetOptions){
					if(!hasSourceConflit(root, option)){
						possibleOptions.add(option);
					}
				}
				LinkedList<TranslationHypothesis> newHypothesis = generateNewHypothesis(root, possibleOptions);
				stack.addAll(newHypothesis);
			}
		}
		return finalHypothesis;
	}
	
	protected static boolean hasSourceConflit(TranslationHypothesis hyp, TranslationOption option){
		for(int i = option.getSourceS(); i<=option.getSourceE(); i++){
			if(hyp.getUsedSourceWords()[i]){
				return true;
			}
		}
		return false;
	}
	
	protected LinkedList<TranslationHypothesis> generateNewHypothesis(TranslationHypothesis root, LinkedList<TranslationOption> options){
		LinkedList<TranslationHypothesis> newHyps = new LinkedList<TranslationHypothesis>();
		for(TranslationOption option : options){
			TranslationHypothesis newHyp = root.addOption(option);
			for(TranslationHypothesisScore scorer : scores.values()){
				scorer.addOption(root, newHyp, option);
			}
			newHyps.add(newHyp);
		}
		return newHyps;
	}
	
	protected boolean isHypothesisComplete(TranslationHypothesis hyp){
		if(hyp.translatedWordIndex == targetSize-1){
			for(boolean b : hyp.usedSourceWords){
				if(!b){
					return false;
				}
			}
			return true;
		}
		return false;
	}
	
	public void addOptionScore(String id, TranslationHypothesisScore score){
		scores.put(id, score);
	}
	
	public TranslationHypothesisScore getOptionScore(String id){
		return scores.get(id);
	}
	
	public static void main(String[] args) throws Exception{
		LinkedList<TranslationOption> options = new LinkedList<TranslationOption>();
		options.add(new TranslationOption(null, 0,1,0,0));
		options.add(new TranslationOption(null, 1,1,0,0));
		options.add(new TranslationOption(null, 1,2,0,0));
		options.add(new TranslationOption(null, 0,2,0,0));
		options.add(new TranslationOption(null, 2,3,1,1));
		options.add(new TranslationOption(null, 3,3,1,1));
		options.add(new TranslationOption(null, 1,3,0,1));
		options.add(new TranslationOption(null, 0,3,0,1));
		TranslationDecoder decoder = new TranslationDecoder(options, 4, 2);
		NumberOfOptionsScore optionsScore = new NumberOfOptionsScore();
		decoder.addOptionScore("options", optionsScore);
		LinkedList<TranslationHypothesis> finalHpys = decoder.getTranslationHypothesis();
		for(TranslationHypothesis hyp : finalHpys){
			System.err.println(optionsScore.getNumberOfOptionsForHpy(hyp));			
		}
	}
}